from abc import abstractmethod
from typing import Dict, Union

from requests import Response



# 定义流式响应包装类
class StreamingResponse:
    def __init__(self, data: bytes, original_resp: Response):
        self.data = data
        self.headers = dict(original_resp.headers)
        self.headers.pop('Content-Length', None)  # 移除长度头（流式传输）
        self.status_code = original_resp.status_code

    def iter_content(self, chunk_size=1024):
        # 模拟流式响应
        bytes_data = self.data
        for i in range(0, len(bytes_data), chunk_size):
            yield bytes_data[i:i + chunk_size]


class ResponseHandler:
    @abstractmethod
    def handle(self, host: str, response: Response) -> Union[bytes, str, Response]:
        pass
    
    @staticmethod
    def is_text(response: Response) -> bool:
        """判断响应是否为文本类型
        
        Args:
            response: requests.Response对象
            
        Returns:
            bool: 如果是文本类型返回True，否则返回False
        """
        if 'content-type' not in response.headers:
            print('Content-Type not found.')
            return False
            
        content_type = response.headers['content-type'].lower()
        text_types = [
            'text/',
            'application/json',
            'application/xml',
            'application/javascript',
            'application/x-javascript',
            'application/x-www-form-urlencoded'
        ]
        
        _is_text = any(content_type.startswith(t) for t in text_types)
        print('is text: ' + str(_is_text))
        return _is_text




class ResponseManager:
    handlers: Dict[str, ResponseHandler] = {}

    @classmethod
    def register(cls, proxy_path: str, handler: ResponseHandler):
        if proxy_path in cls.handlers:
            raise ValueError(f"Handler for {proxy_path} already registered")
        print(f'Register response handler for {proxy_path}')
        cls.handlers[proxy_path] = handler

    @classmethod
    def handle(cls, proxy_path: str, host: str, response: Response) -> Union[StreamingResponse, bool]:
        if proxy_path not in cls.handlers:
            return False
        result = cls.handlers[proxy_path].handle(host, response)

        if isinstance(response, Response):
            return response

        if isinstance(result, str):
            result = result.encode('utf-8')

        return StreamingResponse(result, response)
